from typing import TYPE_CHECKING, Union, overload

from zkl_aiutils_training import Resumable
from zkl_pyutils_fsspec import resolve_child_fs, resolve_local_path

from .metrics import MetricOrValue, get_metric_scalar_value
from .metrics_logger import MetricsLogger
from .resumable_fs import FsResumeArgs

if TYPE_CHECKING:
    from torch.utils.tensorboard import SummaryWriter


class TensorBoardMetricsLogger(MetricsLogger, Resumable):
    @overload
    def __init__(self, writer: 'SummaryWriter', *, step_metric_name: str = "step"):
        ...

    @overload
    def __init__(self, summaries_dir_path: str, *, step_metric_name: str = "step"):
        ...

    @overload
    def __init__(self, *, step_metric_name: str = "step"):
        ...

    def __init__(self,
        arg: Union['SummaryWriter', str, None] = None, *,
        writer: 'SummaryWriter' = None,
        summaries_dir_path: str = None,
        step_metric_name: str = "step",
    ):
        if writer is None:
            from torch.utils.tensorboard import SummaryWriter
            if arg is not None:
                writer = arg if isinstance(arg, SummaryWriter) else SummaryWriter(arg)
            elif summaries_dir_path is not None:
                writer = SummaryWriter(summaries_dir_path)

        self.writer = writer
        self.step_metric_name = step_metric_name

    def on_resume(self, args: FsResumeArgs):
        if self.writer is None:
            try:
                summaries_dir_path = resolve_local_path(resolve_child_fs(args.fs, "summaries"))
            except TypeError:
                summaries_dir_path = None

            if summaries_dir_path is not None:
                from torch.utils.tensorboard import SummaryWriter
                self.writer = SummaryWriter(summaries_dir_path)

    def log(self, metrics: dict[str, MetricOrValue]):
        if self.writer is None:
            return

        metrics_value = {}
        for name, value in metrics.items():
            value = get_metric_scalar_value(value)
            if value is not None:
                metrics_value[name] = value
        step = metrics_value.pop("train/" + self.step_metric_name)

        for name, value in metrics_value.items():
            self.writer.add_scalar(name, value, step)
